/**
 * 
 */
package edu.umd.clip.lm.util.algo;

import org.apache.commons.math.util.OpenIntToDoubleHashMap;

import edu.umd.clip.jobs.Job;
import edu.umd.clip.jobs.JobGroup;
import edu.umd.clip.jobs.JobManager;
import edu.umd.clip.lm.util.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class FastMDI<S, T> {
	private final S[] wordVocab;
	private final T[] otherVocab;
	
	private CompactReadOnlyInt2DoubleHashMap[] probabilities;
	private OpenIntToDoubleHashMap[] initProbabilities;
	
	private static final int MAX_ITERATIONS = 15;
	private static final double SORTA_INFINITY = 10;
	private JobGroup jobGroup;
	private Notifier notifier;
	
	public FastMDI(S[] wordVocab, T[] otherVocab) {
		this.wordVocab = wordVocab;
		this.otherVocab = otherVocab;
		
		this.initProbabilities = new OpenIntToDoubleHashMap[wordVocab.length];
	}
	
	public void setProb(int wordIdx, int otherIdx, double prob) {
		OpenIntToDoubleHashMap map = initProbabilities[wordIdx];
		if (map == null) {
			map = new OpenIntToDoubleHashMap(0.0);
			initProbabilities[wordIdx] = map;
		}
		map.put(otherIdx, prob);
	}
	
	public void addProb(int wordIdx, int otherIdx, double prob) {
		OpenIntToDoubleHashMap map = initProbabilities[wordIdx];
		if (map == null) {
			map = new OpenIntToDoubleHashMap(0.0);
			initProbabilities[wordIdx] = map;
		}
		double oldProb = map.put(otherIdx, prob);
		if (oldProb != 0.0) {
			map.put(otherIdx, oldProb + prob); 
		}
	}
	
	private synchronized void finishInitialization() {
		if (probabilities == null) {
			JobManager manager = JobManager.getInstance();
			JobGroup group = manager.createJobGroup("MDI: init");
			
			probabilities = new CompactReadOnlyInt2DoubleHashMap[wordVocab.length];
			
			final CompactReadOnlyInt2DoubleHashMap uniformDist;
			{
				int keys[] = new int[otherVocab.length];
				double values[] = new double[otherVocab.length];

				double prob = 1.0 / otherVocab.length;
				
				for(int i=0; i<otherVocab.length; ++i) {
					keys[i] = i;
					values[i] = prob;
				}
				uniformDist = new CompactReadOnlyInt2DoubleHashMap(keys, values);
			}
			
			for(int i=0; i<wordVocab.length; ++i) {
				final int wordIdx = i;
				
				Runnable run = new Runnable() {
					@Override
					public void run() {
						OpenIntToDoubleHashMap initHash = initProbabilities[wordIdx];
						if (initHash == null) {
							probabilities[wordIdx] = uniformDist;
							return;
						}
						int size = initHash.size();
						
						int keys[] = new int[size];
						double probs[] = new double[size];

						int pos = 0;
						double totalProb = 0;
						for(OpenIntToDoubleHashMap.Iterator it = initHash.iterator(); it.hasNext();) {
							it.advance();
							
							keys[pos] = it.key();
							double prob = it.value();
							probs[pos] = prob;
							totalProb += prob;
							++pos;
						}
						if (totalProb != 1.0) {
							for(int i=0; i<probs.length; ++i) {
								probs[i] /= totalProb;
							}
						}
						
						CompactReadOnlyInt2DoubleHashMap hash = new CompactReadOnlyInt2DoubleHashMap(keys, probs, 0.0);
						probabilities[wordIdx] = hash;
						
						initProbabilities[wordIdx] = null;
					}
				};
				Job job = new Job(run, "a job");
				manager.addJob(group, job);
			}
			group.join();
			
			initProbabilities = null;
		}
	}
	
	private double getProb(int wordIdx, int otherIdx) {
		return probabilities[wordIdx].get(otherIdx);
	}
	
	private double[] getCenter(int words[]) {
		return getCenter(words, 0, words.length);
	}
	
	private double[] getCenter(int words[], int offset, int length) {
		double center[] = new double[otherVocab.length];
		
		for(int idx=offset; idx<offset+length; ++idx) {
			int word = words[idx];
			CompactReadOnlyInt2DoubleHashMap wordProbs = probabilities[word];
			int keys[] = wordProbs.keys();
			double probs[] = wordProbs.values();
			for(int i=0; i<keys.length; ++i) {
				center[keys[i]] += probs[i];
			}
		}
		
		for(int i=0; i<center.length; ++i) {
			center[i] /= words.length;
		}
		return center;
	}
	
	private double distance(int wordIdx1, int wordIdx2) {
		// KL-divergence
		double d = 0;
		
		CompactReadOnlyInt2DoubleHashMap probs1 = probabilities[wordIdx1];
		CompactReadOnlyInt2DoubleHashMap probs2 = probabilities[wordIdx2];

		int keys[] = probs1.keys();
		double probs[] = probs1.values();
		
		for(int i=0; i<keys.length; ++i) {
			int otherIdx = keys[i];
			double p1 = probs[i];
			if (p1 > 0) {
				double p2 = probs2.get(otherIdx);
				if (p2 > 0) {
					d += p1 * Math.log(p1 / p2);
				} else {
					d+= SORTA_INFINITY;
				}
			}
		}
		
		return d;
	}
	
	private double distanceToCenter(int wordIdx, double[] center) {
		double d = 0;

		CompactReadOnlyInt2DoubleHashMap probs = probabilities[wordIdx];
		
		int keys[] = probs.keys();
		double wordProbs[] = probs.values();
		
		for(int i=0; i<keys.length; ++i) {
			int otherIdx = keys[i];
			double wordProb = wordProbs[i];
			
			double centerProb = center[otherIdx];
			
			if (centerProb > 0) {
				d += wordProb * Math.log(wordProb / centerProb);
			} else {
				d += SORTA_INFINITY;
			}
		}
		
		if (d < 0.0) {
			System.err.printf("Negative distance: %f\n", d);
			return 0.0;
		}
		return d;
	}

	private void partition(final int[] words, final double[] center) {
		Runnable runnable = new Runnable() {
			public void run() {
				partitionImpl(words, center);
			}
		};
		Job job = new Job(runnable, "MDI");
		JobManager manager = JobManager.getInstance();
		manager.addJob(jobGroup, job);
	}
	

	private void partitionImpl(int[] words, double[] center) {
		final int HALF_SIZE = words.length / 2 + 1;
		
		// find two points closest to the center
		double firstVal = Double.POSITIVE_INFINITY;
		double secondVal = Double.POSITIVE_INFINITY;
		int firstWord = -1;
		int secondWord = -1;
		
		GrowingIntArray cluster1 = new GrowingIntArray(HALF_SIZE);
		GrowingIntArray cluster2 = new GrowingIntArray(HALF_SIZE);
		
		double[] newCenter1 = null;
		double[] newCenter2 = null;

		if (words.length > 2) {
			for(int word : words) {
				double dist = distanceToCenter(word, center);
				if (dist < secondVal) {
					if (dist < firstVal) {
						secondVal = firstVal;
						secondWord = firstWord;
						firstVal = dist;
						firstWord = word;
					} else {
						secondVal = dist;
						secondWord = word;
					}
				}
			}
			if (distance(firstWord, secondWord) <= 0.0) {
				//System.out.printf("zero distance between '%d' and '%d'\n", firstWord, secondWord);
				// pick an arbitrary second word
				for(int word : words) {
					if (word == firstWord || word == secondWord) continue;
					if (distance(firstWord, word) > 0.0) {
						secondWord = word;
						break;
					}
				}
				//System.out.printf("selected '%s' and '%s'\n", firstWord, secondWord);
			}
			
			double[] center1 = getCenter(new int[]{firstWord});
			double[] center2 = getCenter(new int[]{secondWord});
			for(int word : words) {
				double dist1 = distanceToCenter(word, center1);
				double dist2 = distanceToCenter(word, center2);
				if (dist1 < dist2) {
					cluster1.add(word);
				} else {
					cluster2.add(word);
				}
			}
	
			int changed = 1;
			int iteration = 0;
			while(++iteration <= MAX_ITERATIONS && changed > 0) {
				//System.out.printf("Iteration %d/%d\n", iteration, MAX_ITERATIONS);
				changed = 0;
				newCenter1 = getCenter(cluster1.array(), 0, cluster1.size());
				newCenter2 = getCenter(cluster2.array(), 0, cluster2.size());
				
				GrowingIntArray tmpCluster1 = new GrowingIntArray(HALF_SIZE);
				GrowingIntArray tmpCluster2 = new GrowingIntArray(HALF_SIZE);
				
				int cluster1array[] = cluster1.array();
				for(int i=0; i<cluster1.size(); ++i) {
					int word = cluster1array[i];
					double dist1 = distanceToCenter(word, newCenter1);
					double dist2 = distanceToCenter(word, newCenter2);
					if (dist1 <= dist2) {
						tmpCluster1.add(word);
					} else {
						tmpCluster2.add(word);
						changed++;
						//System.out.printf("changing %s : %g > %g\n", word, dist1, dist2);
					}
				}
				
				int cluster2array[] = cluster2.array();
				for(int i=0; i<cluster2.size(); ++i) {
					int word = cluster2array[i];
					double dist1 = distanceToCenter(word, newCenter1);
					double dist2 = distanceToCenter(word, newCenter2);
					if (dist2 <= dist1) {
						tmpCluster2.add(word);
					} else {
						tmpCluster1.add(word);
						changed++;
						//System.out.printf("changing %s : %g < %g\n", word, dist1, dist2);
					}
				}
				cluster1 = tmpCluster1;
				cluster2 = tmpCluster2;
				//System.out.printf("changed: %d/%d\n", changed, vocab.size());
				if (cluster1.size() == 0 || cluster2.size() == 0) {
					GrowingIntArray cl = cluster1.size() == 0 ? cluster2 : cluster1;
					//System.out.printf("Cannot split %s, split them randomly\n", cl.toString());
					tmpCluster1 = new GrowingIntArray(HALF_SIZE);
					tmpCluster2 = new GrowingIntArray(HALF_SIZE);
					int clusterArray[] = cl.array();
					boolean b = true;
					for(int i=0; i<cl.size(); ++i) {
						int word = clusterArray[i];
						if (b) {
							tmpCluster1.add(word);
						} else {
							tmpCluster2.add(word);
						}
						b = !b;
					}
					cluster1 = tmpCluster1;
					cluster2 = tmpCluster2;
				}
			}
		} else if (words.length == 2) {
			int word1 = words[0];
			int word2 = words[1];
			cluster1.add(word1);
			cluster2.add(word2);
			newCenter1 = getCenter(new int[]{word1});
			newCenter2 = getCenter(new int[]{word2});
		} else {
			return;			
		}
		
		cluster1.trim();
		cluster2.trim();
		
		if (notifyNewCluster(words, cluster1.array(), cluster2.array())) {
			partition(cluster1.array(), newCenter1);
			partition(cluster2.array(), newCenter2);
		}
	}
	private boolean notifyNewCluster(int[] words, int[] cluster1, int[] cluster2) {
		if (notifier != null) {
			return notifier.notify(words, cluster1, cluster2);
		}
		return false;
	}

	public void partition(int[] words) {
		finishInitialization();
		
		jobGroup = JobManager.getInstance().createJobGroup("MDI");
		
		partition(words, getCenter(words));
		
		jobGroup.join();
	}
	
	public static interface Notifier {
		public boolean notify(int[] words, int[] cluster1, int[] cluster2);
	}

	public Notifier getNotifier() {
		return notifier;
	}

	public void setNotifier(Notifier notifier) {
		this.notifier = notifier;
	}
}
